"""特征工程模块：技术指标计算"""
import pandas as pd
import numpy as np
from typing import List, Optional


def calculate_sma(data: pd.Series, window: int) -> pd.Series:
    """
    计算简单移动平均线（SMA）
    
    Args:
        data: 价格序列
        window: 窗口大小
        
    Returns:
        SMA序列
    """
    return data.rolling(window=window).mean()


def calculate_ema(data: pd.Series, window: int) -> pd.Series:
    """
    计算指数移动平均线（EMA）
    
    Args:
        data: 价格序列
        window: 窗口大小
        
    Returns:
        EMA序列
    """
    return data.ewm(span=window, adjust=False).mean()


def calculate_macd(
    data: pd.Series,
    fast_period: int = 12,
    slow_period: int = 26,
    signal_period: int = 9
) -> pd.DataFrame:
    """
    计算MACD指标
    
    Args:
        data: 收盘价序列
        fast_period: 快线周期
        slow_period: 慢线周期
        signal_period: 信号线周期
        
    Returns:
        包含MACD、信号线、柱状图的DataFrame
    """
    ema_fast = calculate_ema(data, fast_period)
    ema_slow = calculate_ema(data, slow_period)
    macd_line = ema_fast - ema_slow
    signal_line = calculate_ema(macd_line, signal_period)
    histogram = macd_line - signal_line
    
    return pd.DataFrame({
        'MACD': macd_line,
        'MACD_signal': signal_line,
        'MACD_hist': histogram
    })


def calculate_rsi(data: pd.Series, period: int = 14) -> pd.Series:
    """
    计算相对强弱指标（RSI）
    
    Args:
        data: 价格序列
        period: 计算周期，默认14
        
    Returns:
        RSI序列
    """
    delta = data.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
    
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    
    return rsi


def calculate_atr(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    period: int = 14
) -> pd.Series:
    """
    计算平均真实波幅（ATR）
    
    Args:
        high: 最高价序列
        low: 最低价序列
        close: 收盘价序列
        period: 计算周期，默认14
        
    Returns:
        ATR序列
    """
    high_low = high - low
    high_close = np.abs(high - close.shift())
    low_close = np.abs(low - close.shift())
    
    ranges = pd.concat([high_low, high_close, low_close], axis=1)
    true_range = ranges.max(axis=1)
    
    atr = true_range.rolling(window=period).mean()
    
    return atr


def calculate_bollinger_bands(
    data: pd.Series,
    window: int = 20,
    num_std: float = 2.0
) -> pd.DataFrame:
    """
    计算布林带（Bollinger Bands）
    
    Args:
        data: 价格序列
        window: 移动平均窗口
        num_std: 标准差倍数
        
    Returns:
        包含上轨、中轨、下轨的DataFrame
    """
    sma = calculate_sma(data, window)
    std = data.rolling(window=window).std()
    
    upper_band = sma + (std * num_std)
    lower_band = sma - (std * num_std)
    
    return pd.DataFrame({
        'BB_upper': upper_band,
        'BB_middle': sma,
        'BB_lower': lower_band
    })


def calculate_momentum(data: pd.Series, period: int = 10) -> pd.Series:
    """
    计算动量指标
    
    Args:
        data: 价格序列
        period: 计算周期
        
    Returns:
        动量序列
    """
    return data.pct_change(periods=period)


def calculate_volume_features(volume: pd.Series, close: pd.Series) -> pd.DataFrame:
    """
    计算成交量相关特征
    
    Args:
        volume: 成交量序列
        close: 收盘价序列
        
    Returns:
        包含成交量特征的DataFrame
    """
    # 成交量移动平均
    volume_ma_5 = volume.rolling(window=5).mean()
    volume_ma_20 = volume.rolling(window=20).mean()
    
    # 成交量比率
    volume_ratio = volume / volume_ma_20
    
    # 量价关系：价格变化与成交量的相关性
    price_change = close.pct_change()
    volume_change = volume.pct_change()
    
    # 价量相关性（滚动窗口）
    price_volume_corr = price_change.rolling(window=20).corr(volume_change)
    
    return pd.DataFrame({
        'volume_ma5': volume_ma_5,
        'volume_ma20': volume_ma_20,
        'volume_ratio': volume_ratio,
        'volume_change': volume_change,
        'price_volume_corr': price_volume_corr
    })


def calculate_volatility_features(close: pd.Series, windows: List[int] = [5, 10, 20]) -> pd.DataFrame:
    """
    计算波动率特征
    
    Args:
        close: 收盘价序列
        windows: 计算窗口列表
        
    Returns:
        包含波动率特征的DataFrame
    """
    returns = close.pct_change()
    
    features = {}
    for window in windows:
        features[f'volatility_{window}'] = returns.rolling(window=window).std() * np.sqrt(252)
    
    return pd.DataFrame(features)


def create_features(
    df: pd.DataFrame,
    price_col: str = '收盘价',
    high_col: str = '最高价',
    low_col: str = '最低价',
    close_col: str = '收盘价',
    volume_col: str = '成交量',
    include_indicators: Optional[List[str]] = None
) -> pd.DataFrame:
    """
    创建完整的特征集
    
    Args:
        df: 包含K线数据的DataFrame
        price_col: 价格列名（用于SMA/EMA，默认使用收盘价）
        high_col: 最高价列名
        low_col: 最低价列名
        close_col: 收盘价列名
        volume_col: 成交量列名
        include_indicators: 要包含的技术指标列表，None表示包含所有
        
    Returns:
        包含所有特征的DataFrame
    """
    if include_indicators is None:
        include_indicators = ['sma', 'ema', 'macd', 'rsi', 'atr', 'bollinger', 'momentum', 'volume', 'volatility']
    
    features_list = []
    
    # 确保使用正确的列名
    close = df[close_col] if close_col in df.columns else df['收盘价']
    high = df[high_col] if high_col in df.columns else df['最高价']
    low = df[low_col] if low_col in df.columns else df['最低价']
    volume = df[volume_col] if volume_col in df.columns else df['成交量']
    price = df[price_col] if price_col in df.columns else close
    
    # SMA特征
    if 'sma' in include_indicators:
        features_list.append(calculate_sma(price, 5).rename('SMA_5'))
        features_list.append(calculate_sma(price, 10).rename('SMA_10'))
        features_list.append(calculate_sma(price, 20).rename('SMA_20'))
        features_list.append(calculate_sma(price, 60).rename('SMA_60'))
    
    # EMA特征
    if 'ema' in include_indicators:
        features_list.append(calculate_ema(price, 5).rename('EMA_5'))
        features_list.append(calculate_ema(price, 10).rename('EMA_10'))
        features_list.append(calculate_ema(price, 20).rename('EMA_20'))
    
    # MACD特征
    if 'macd' in include_indicators:
        macd_features = calculate_macd(close)
        features_list.append(macd_features['MACD'].rename('MACD'))
        features_list.append(macd_features['MACD_signal'].rename('MACD_signal'))
        features_list.append(macd_features['MACD_hist'].rename('MACD_hist'))
    
    # RSI特征
    if 'rsi' in include_indicators:
        features_list.append(calculate_rsi(close, 14).rename('RSI_14'))
        features_list.append(calculate_rsi(close, 9).rename('RSI_9'))
    
    # ATR特征
    if 'atr' in include_indicators:
        features_list.append(calculate_atr(high, low, close, 14).rename('ATR_14'))
    
    # 布林带特征
    if 'bollinger' in include_indicators:
        bb_features = calculate_bollinger_bands(close, 20, 2.0)
        features_list.append(bb_features['BB_upper'].rename('BB_upper'))
        features_list.append(bb_features['BB_middle'].rename('BB_middle'))
        features_list.append(bb_features['BB_lower'].rename('BB_lower'))
        # 价格相对于布林带的位置
        features_list.append(((close - bb_features['BB_lower']) / (bb_features['BB_upper'] - bb_features['BB_lower'])).rename('BB_position'))
    
    # 动量特征
    if 'momentum' in include_indicators:
        features_list.append(calculate_momentum(close, 5).rename('momentum_5'))
        features_list.append(calculate_momentum(close, 10).rename('momentum_10'))
        features_list.append(calculate_momentum(close, 20).rename('momentum_20'))
    
    # 成交量特征
    if 'volume' in include_indicators:
        volume_features = calculate_volume_features(volume, close)
        for col in volume_features.columns:
            features_list.append(volume_features[col].rename(col))
    
    # 波动率特征
    if 'volatility' in include_indicators:
        volatility_features = calculate_volatility_features(close, [5, 10, 20])
        for col in volatility_features.columns:
            features_list.append(volatility_features[col].rename(col))
    
    # 价格变化特征
    features_list.append(close.pct_change(1).rename('returns_1'))
    features_list.append(close.pct_change(5).rename('returns_5'))
    features_list.append(close.pct_change(10).rename('returns_10'))
    
    # 高低价特征
    features_list.append(((high - low) / close).rename('price_range'))
    features_list.append(((close - low) / (high - low)).rename('price_position'))
    
    # 合并所有特征
    features_df = pd.concat(features_list, axis=1)
    
    # 填充NaN值
    features_df = features_df.bfill().fillna(0)
    
    return features_df


def create_target(
    df: pd.DataFrame,
    close_col: str = '收盘价',
    method: str = 'classification',
    future_periods: int = 1,
    threshold: float = 0.0
) -> pd.Series:
    """
    创建目标变量
    
    Args:
        df: 包含K线数据的DataFrame
        close_col: 收盘价列名
        method: 方法类型，'classification'（分类）或 'regression'（回归）
        future_periods: 未来预测周期数
        threshold: 分类阈值（仅分类任务使用）
        
    Returns:
        目标变量Series
    """
    close = df[close_col] if close_col in df.columns else df['收盘价']
    
    if method == 'classification':
        # 分类：未来N期涨跌（1=上涨，0=下跌）
        future_return = close.shift(-future_periods) / close - 1
        target = (future_return > threshold).astype(int)
    else:
        # 回归：未来N期收益率
        target = (close.shift(-future_periods) / close - 1)
    
    return target

